package com.leo.boot.scheduler.util;

import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Arrays;
import java.util.concurrent.atomic.AtomicInteger;

@UtilityClass
@Slf4j
public class TaskDecorators {

    AtomicInteger counter = new AtomicInteger();

    public static TaskDecorator reduce(TaskDecorator... taskDecorators) {
        return Arrays.stream(taskDecorators)
                .reduce((a, b) -> runnable -> a.decorate(b.decorate(runnable)))
                .orElse(nothing());
    }

    public static TaskDecorator nothing() {
        return runnable -> runnable;
    }

    public static TaskDecorator stat() {
        return runnable -> () -> {
            int taskId = counter.incrementAndGet();
            long start = System.currentTimeMillis();
            log.debug("Task {} begin", taskId);
            try {
                runnable.run();
            } finally {
                log.debug("Task {} end, total time {} ms", taskId, stat());
            }
        };
    }

    public static TaskDecorator trace() {
        return runnable -> () -> {
            String tid = MDC.get("tid");
            try {
                MDC.put("tid", tid);
                runnable.run();
            } finally {
                MDC.remove("tid");
            }
        };
    }
}

